import os
import csv
import json
import shutil
import cv2

train_datas = []
csvFile = open("./data/data_split.csv", "r")
reader = csv.reader(csvFile)
for item in reader:
    if reader.line_num == 1:  # 忽略第一行
        continue
    if item[-1] == "validation":  # validation training
        train_datas.append(item[0])
file_train_txt = open('./data/custom/valid.txt', 'w')
for train_data in train_datas:
    print(train_data)
    file = open('./data/json/' + train_data + '.json')
    dict = json.load(file)
    file.close()
    for key, values in dict.items():
        # train.txt
        file_train_txt.write('./data/custom/images/' + train_data + '_' + key + '.jpg' + '\n')
        # images
        shutil.copy('/home/he/files/baoweichao/yolov3_helmat/data/all_data/' + train_data + '/' + key + '.jpg',
                    './data/custom/images/' + train_data + '_' + key + '.jpg')
        # labels
        label_file = open('./data/custom/labels/' + train_data + '_' + key + '.txt', 'w')
        for value in values:
            label_file.write(str(int(value[0])) + ' ' + str(value[1] / 1920.0) + ' ' +
                             str((value[2] + 420) / 1920.0) + ' ' + str(value[3] / 1920.0) +
                             ' ' + str(value[4] / 1920.0) + '\n')
        label_file.close()
file_train_txt.close()
